`include "top.vh"

module mem_stage (
    input                          clk           ,
    input                          reset         ,
    //allowin
    input                          ws_allowin    ,
    output                         ms_allowin    ,
    //from es
    input                          es_to_ms_valid,
    input  [`es_2_ms_bus_wid - 1:0] es_to_ms_bus  ,
    //to ws
    output                         ms_to_ws_valid,
    output [`ms_2_ws_bus_wid - 1:0] ms_to_ws_bus  ,
    //to id
    output [37:0] ms_to_id_bus ,
    //from data-sram
    input  [31:0] data_sram_rdata,
    //csr data block
    output read_csr,
    //ms_ex
    output ms_ex_out,
    output ms_ertn_flush,
    //ws_ex
    input  ws_ex,
    input  ws_ertn_flush
);

reg         ms_valid;
wire        ms_ready_go;

wire ms_res_from_mem;
wire [ 2:0] ms_ld_op;
wire [31:0] ms_alu_result;
wire ms_gr_we;
wire [ 4:0] ms_dest;
wire [31:0] ms_pc;
wire ms_rf_we;
wire [31:0] ms_final_result;
wire [ 7:0] load_byte;
wire [15:0] load_half;
wire [31:0] ms_result;
wire [ 1:0] ms_vaddr2;

wire csr_re;
wire [`CSR_SIG_wid-1:0] csr_signal;
wire es_ex, es_ertn;
wire ms_ex, ms_ertn;
wire [`EX_SIG_wid-1:0] es_ex_signal, ms_ex_signal;

assign ms_ex = es_ex;
assign ms_ex_signal =   es_ex ? es_ex_signal :
                        4'h0;
assign ms_ertn = es_ertn;

assign csr_re = csr_signal[`CSR_SIG_RE];
assign read_csr = csr_re && ms_valid;
assign ms_ex_out = ms_ex && ms_valid;
assign ms_ertn_flush = ms_ertn && ms_valid;

reg [`es_2_ms_bus_wid - 1:0] es_to_ms_bus_reg;

always @(posedge clk ) begin
    if (ms_allowin) begin
        es_to_ms_bus_reg  <= es_to_ms_bus;
    end
end

assign {
        csr_signal,     //e81
        es_ex,             //1
        es_ex_signal,   //16
        es_ertn,           //1
            
        ms_ld_op,
        ms_vaddr2,
        ms_alu_result,
        ms_gr_we,
        ms_dest,
        ms_pc
} = es_to_ms_bus_reg;

assign ms_res_from_mem = ms_ld_op != 3'b0 && !ms_ex;

assign ms_to_ws_bus = {
                        csr_signal,     //81
                        ms_ex,             //1
                        ms_ex_signal,   //16
                        ms_ertn,           //1
            
                        ms_gr_we,
                        ms_dest,
                        ms_final_result,
                        ms_pc
                    };
assign ms_to_id_bus = {

                        ms_final_result,//37:6
                        ms_rf_we,//5
                        ms_dest//0:4
};

assign ms_rf_we = ms_gr_we & ms_valid;
assign ms_ready_go = 1'b1;
assign ms_allowin = (~ms_valid) | (ms_ready_go & ws_allowin);
assign ms_to_ws_valid = ms_valid & ms_ready_go;

always @(posedge clk) begin
    if (reset) begin
        ms_valid <= 1'b0;
    end
    else if (ws_ex || ws_ertn_flush) 
        ms_valid <= 1'b0;
    else if (ms_allowin) begin
        ms_valid <= es_to_ms_valid;
    end
end

assign load_byte   =    {8 {ms_vaddr2 == 2'b11}} & data_sram_rdata[31:24] |
                        {8 {ms_vaddr2 == 2'b10}} & data_sram_rdata[23:16] |
                        {8 {ms_vaddr2 == 2'b01}} & data_sram_rdata[15: 8] |
                        {8 {ms_vaddr2 == 2'b00}} & data_sram_rdata[ 7: 0];
assign load_half   =    {16{ms_vaddr2[1]}}       & data_sram_rdata[31:16] |
                        {16{~ms_vaddr2[1]}}      & data_sram_rdata[15: 0];

assign ms_result   =    ms_ld_op == 3'b111 ?    data_sram_rdata :                 //ld.w
                        ms_ld_op == 3'b101 ?    {{24{load_byte[7]}}, load_byte} : //ld.b
                        ms_ld_op == 3'b001 ?    {24'b0, load_byte} :              //ld.bu
                        ms_ld_op == 3'b110 ?    {{16{load_half[15]}}, load_half} ://ld.h
                        ms_ld_op == 3'b010 ?    {16'b0 , load_half} :             //ld.hu
                                                data_sram_rdata;
assign ms_final_result = ms_res_from_mem ? ms_result : ms_alu_result;
endmodule